#!/usr/bin/python
# -*- coding: utf8 -*-

from enthought.traits.api import HasTraits, Instance
from enthought.traits.ui.api import View, Item
from enthought.chaco.api import HPlotContainer, VPlotContainer, ArrayPlotData, Plot, PlotAxis, PlotLabel, \
                                create_line_plot, OverlayPlotContainer, add_default_axes, add_default_grids, Legend
from enthought.enable.component_editor import ComponentEditor
from numpy import linspace, array

auto_colors = ['red','gold','darkblue','green','purple','brown']

def color(i):
    '''Devuelve un color para el elemento i de una secuencia de gráficos'''
    return auto_colors[i % len(auto_colors)]

def derivar(arr):
    '''Calcula una aproximación de la derivada de una secuencia de números'''
    if len(arr) <= 1: return [0] * len(arr)

    tot = len(arr)
    res = [None]
    for k in range(tot-1):
        res.append((-arr[k] + arr[k+1]))
    res[0] = res[1]
    res[-1] = res[-2]
    return res


class PlotterEnfermedad(HasTraits):
    plot = Instance(HPlotContainer)
    traits_view = View(Item('plot', editor=ComponentEditor(), show_label=False),
                       width=1200, height=950, resizable=True, title="Evolución de la enfermedad")
    
    def __init__(self, valores):
        super(PlotterEnfermedad, self).__init__()

        # Chequeos
        assert(len(valores) > 0)

        cuenta = None
        for key, value in valores.iteritems():
            if cuenta is None:
                cuenta = len(value)
            else:
                assert cuenta == len(value)
        
        # Construcción de los gráficos
        plots = []
        dict_plots = {}
        x = linspace(1, cuenta, cuenta)
        
        i = 0
        for key, value in valores.iteritems():
            c = color(i)
            i = i + 1

            y = array(value)
            ders = array(derivar(y))
        
            # Este container permite dibujar 2 gráficos sobre el mismo plano
            container = OverlayPlotContainer(padding = 26, fill_padding = True,
                                             bgcolor = "lightgray", use_backbuffer=True)
            
            # Gráfico de valores
            l1 = create_line_plot((x, y), color=c, dash="solid")
            add_default_axes(l1)
            add_default_grids(l1)
            l1.bgcolor = 'white'
            l1.index.sort_order = "ascending"
            dict_plots[key] = l1
            container.add(l1)
            
            # Gráfico de la derivada
            l2 = create_line_plot((x, ders), color=c, dash="dash")
            axis = PlotAxis(l2, orientation="right")
            l2.underlays.append(axis)
            container.add(l2)

            # Leyenda
            legend = Legend(component=container, padding=5, align="ul")
            container.overlays.append(legend)
            legend.plots = {key: l1, u'Tasa de variación': l2}

            plots.append(container)


        # Construyo el gráfico global a un costado de los locales
        cont = VPlotContainer(*plots)

        plotdata = ArrayPlotData(index=x, **valores)
        line = Plot(plotdata) 
        i = 0
        for key, value in valores.iteritems():
            c = color(i)
            i = i + 1
            line.plot(("index", key), type="line", color=c)
        line.title = u"Evolución global"
        
        legend = Legend(component=line, padding=10, align="ur")
        line.overlays.append(legend)
        legend.plots = dict_plots

        self.plot = HPlotContainer(cont, line, bgcolor="lightgray")


def plotear_enfermedad(datos):
    p = PlotterEnfermedad(datos)
    p.configure_traits()


if __name__ == "__main__":
    plotear_enfermedad({"Sanos":[1,2,3,4,5,4,3,2,1], "Enfermos":[5,4,3,1,1,65,23,5,2], "Inmunes":[2,6,3,6,2,2,25,76,9]})



